import { post } from "../http";
import { AGENT_BACKEND_API_KEY } from "../auth";
import { BaiduConfig, defaultBaiduConfig } from "./config";

// 定义请求体的接口
interface ChatCompletionRequest {
    model: string;
    messages: {
        role: "user" | "assistant";
        content: string;
    }[];
    /**
     * 模型人设，主要用于人设设定
     */
    system: string;
}

interface ChatCompletionResponse {
    id: string;
    object: "chat.completion";
    created: number;
    sentence_id?: number; // 只有在流式接口模式下会返回该字段
    is_end?: boolean; // 只有在流式接口模式下会返回该字段
    is_truncated: boolean;
    finish_reason: "normal" | "stop" | "length" | "content_filter";
    search_info?: {
        // 这里可以根据具体需求定义搜索数据的具体字段
        [key: string]: any;
    }; // 当请求参数enable_citation或enable_trace为true，并且触发搜索时，会返回该字段
    result: string;
    need_clear_history: boolean;
    flag: number;
    ban_round: number;
    usage: {
        prompt_tokens: number;
        completion_tokens: number;
        total_tokens: number;
    };
}

/**
 * 获取聊天完成响应
 * @param request 请求体
 * @returns
 */
export async function getChatCompletion(
    request: ChatCompletionRequest,
    apiKey: string,
    modelConfig: BaiduConfig
): Promise<ChatCompletionResponse> {
    const key = apiKey ?? AGENT_BACKEND_API_KEY;
    if (!key) {
        throw new Error("API key is missing");
    }

    const config = { ...defaultBaiduConfig, ...modelConfig };

    const requestBody = {
        ...config,
        model: request.model,
        messages: request.messages
    };
    key;
    try {
        const response = await post<ChatCompletionResponse>(
            `https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/${(requestBody.model).toLowerCase()}?access_token=${key}`,
            // `https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token=${key}`,
            requestBody
        );

        if (response.status === 200) {
            // 返回助手的回复内容
            return response.data;
        } else {
            console.error(`Unexpected response status: ${response.status}`);
            throw new Error("Failed to get chat completion");
        }
    } catch (error) {
        console.error("Error in getting chat completion:", error);
        throw error;
    }
}
